{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nupic.research.frameworks.dynamic_sparse.networks import GSCHeb\n",
    "from nupic.research.frameworks.dynamic_sparse.models import DSNNWeightedMag\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "network = GSCHeb()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1600])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network.features(torch.randn(1, 1,32,32)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Linear(in_features=1600, out_features=1000, bias=True)\n",
       "  (1): BatchNorm1d(1000, eps=1e-05, momentum=0.1, affine=False, track_running_stats=True)\n",
       "  (2): KWinners(n=1000, percent_on=0.1, boost_strength=1.5, boost_strength_factor=0.9, k_inference_factor=1.5, duty_cycle_period=1000)\n",
       "  (3): Linear(in_features=1000, out_features=12, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network.classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 12])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(4, 1600)\n",
    "x = network.classifier[0](x)\n",
    "x = network.classifier[1](x)\n",
    "x = network.classifier[2](x)\n",
    "x = network.classifier[3](x)\n",
    "\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 12])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network.classifier(torch.randn(4, 1600)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 12])"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network(torch.randn(4, 1,32,32)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "network.init_hebbian()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 12])"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network(torch.randn(4, 1,32,32)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(500):\n",
    "    network(torch.randn(16, 1,32,32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[ 97.,  88., 112.,  ..., 102., 111.,  94.],\n",
       "         [ 96., 137., 102.,  ...,  93., 114.,  99.],\n",
       "         [ 96.,  78., 126.,  ..., 116., 113., 114.],\n",
       "         ...,\n",
       "         [ 61.,  92.,  79.,  ...,  90.,  84.,  79.],\n",
       "         [ 72.,  84.,  70.,  ...,  90., 100.,  70.],\n",
       "         [ 91., 102.,  77.,  ...,  79.,  88.,  78.]]),\n",
       " tensor([[547., 452., 419.,  ..., 336., 554., 174.],\n",
       "         [466., 403., 492.,  ..., 369., 545., 135.],\n",
       "         [438., 506., 384.,  ..., 394., 558., 149.],\n",
       "         ...,\n",
       "         [425., 497., 445.,  ..., 396., 530., 150.],\n",
       "         [521., 439., 418.,  ..., 417., 595., 177.],\n",
       "         [498., 443., 418.,  ..., 385., 547., 159.]])]"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network.coactivations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "network = GSCHeb()\n",
    "model = DSNNWeightedMag(network, dict(on_perc=01.))\n",
    "model.setup()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(4):\n",
    "    network(torch.randn(16, 1,32,32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[0., 0., 0.,  ..., 0., 2., 0.],\n",
       "         [0., 0., 1.,  ..., 0., 1., 1.],\n",
       "         [0., 1., 1.,  ..., 0., 2., 1.],\n",
       "         ...,\n",
       "         [0., 2., 2.,  ..., 0., 1., 0.],\n",
       "         [0., 0., 2.,  ..., 0., 2., 0.],\n",
       "         [0., 1., 1.,  ..., 0., 1., 1.]]),\n",
       " tensor([[5., 5., 2.,  ..., 2., 4., 3.],\n",
       "         [3., 3., 2.,  ..., 4., 5., 3.],\n",
       "         [2., 3., 2.,  ..., 3., 2., 3.],\n",
       "         ...,\n",
       "         [4., 1., 0.,  ..., 1., 2., 1.],\n",
       "         [7., 6., 5.,  ..., 1., 5., 5.],\n",
       "         [7., 4., 2.,  ..., 3., 5., 3.]])]"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network.coactivations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "model._reinitialize_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[0., 0., 0.,  ..., 0., 2., 0.],\n",
       "         [0., 0., 1.,  ..., 0., 1., 1.],\n",
       "         [0., 1., 1.,  ..., 0., 2., 1.],\n",
       "         ...,\n",
       "         [0., 2., 2.,  ..., 0., 1., 0.],\n",
       "         [0., 0., 2.,  ..., 0., 2., 0.],\n",
       "         [0., 1., 1.,  ..., 0., 1., 1.]]),\n",
       " tensor([[5., 5., 2.,  ..., 2., 4., 3.],\n",
       "         [3., 3., 2.,  ..., 4., 5., 3.],\n",
       "         [2., 3., 2.,  ..., 3., 2., 3.],\n",
       "         ...,\n",
       "         [4., 1., 0.,  ..., 1., 2., 1.],\n",
       "         [7., 6., 5.,  ..., 1., 5., 5.],\n",
       "         [7., 4., 2.,  ..., 3., 5., 3.]])]"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "network.coactivations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
